import pandas as pd
import numpy as np


def mark_xl(df, label_window=3, threshold=0.15):
    """
    生成预测三日后涨幅超15%概率的大单特征
    
    参数:
        df: 原始数据 (需含date,code,close,xl,xl_rate等字段)
        label_window: 预测窗口天数 (默认3日)
        threshold: 涨幅阈值 (默认15%)
    
    返回:
        含特征和标签的DataFrame
    """
    df = df.sort_values(['code', 'date']).copy()
    
    # ===== 标签生成 =====
    df['future_close'] = df.groupby('code')['close'].shift(-label_window)
    
    # ===== 大单特征工程 =====
    # 1. 基础特征
    df['xl_net'] = df['xl']
    df['xl_ratio'] = df['xl_rate']

    def _cal_yesterday_feature(df):
        df = df.sort_values(by=['code', 'date'])
        df['pre_xl_ratio'] = df.groupby('code')['xl_ratio'].shift(1)
        return df
    df = df.groupby('code', group_keys=False).apply(_cal_yesterday_feature)
    
    
    # 2. 时间窗口特征
    windows = [3,5,10]
    for w in windows:
        # 滚动窗口计算
        df[f'xl_cumsum_{w}d'] = df.groupby('code')['xl'].transform(lambda s: s.rolling(w).sum())
        df[f'xl_rate_ma_{w}d'] = df.groupby('code')['xl_rate'].transform(lambda s: s.rolling(w).mean())
        df[f'xl_vol_{w}d'] = df.groupby('code')['xl'].transform(lambda s: s.rolling(w).std())
        
    # 3. 量价协同特征
    df['xl_push_eff'] = df['xl'] / (df['turnover'] * df['quote_rate'].abs().replace(0, 0.001))
    df['xl_price_div'] = (df['xl'] > df['xl'].rolling(5).mean()) & (df['close'] < df['close'].rolling(5).mean())
    
    # 4. 异常波动特征
    df['xl_spike'] = (df['xl'] > 2 * df.groupby('code')['xl'].transform(lambda s: s.rolling(20).mean()))
    df['xl_break_high'] = (df['xl'] > df.groupby('code')['xl'].transform(lambda s: s.rolling(10).max()))
    
    # 5. 资金流特征
    df['xl_concentration'] = df['xl'] / df.groupby('code')['xl'].transform('mean')
    df['xl_t_rate'] = df['xl'] / (df['turnover'] / df['close'])
    
    # 6. 持续性特征
    df['xl_positive_days'] = df.groupby('code')['xl'].transform(
        lambda s: s.gt(0).groupby(s.diff().ne(0).cumsum()).cumcount() + 1
    )
    
    # ===== 清理中间列 =====
    df = df.drop(columns=['future_close'])
    #   # 列删除逻辑（根据参数动态调整）
    # base_drop_cols = ['_base']
    # # 安全删除（网页2建议）
    # all_drop_cols = [col for col in (base_drop_cols ) if col in df.columns]
    # df.drop(columns=[''], axis=1, inplace=True, errors='ignore')
    df = drop_base_columns(df)
    return df


def mark_xl_col_growth(
    df: pd.DataFrame,
    target_col: str,
    window: int = 10,
    feature_suffix: Optional[str] = None,
    min_valid_value: float = 1e-6,
    growth_type: str = 'relative',
    trading_col: str = 'is_trading',
    is_drop: bool = True
) -> pd.DataFrame:
    """
    优化版增长率计算 - 保持原有4列输出
    
    改进点：
    1. 零值特殊处理：当基准值为零时，使用替代策略计算增长率
    2. 仅保留最终需要的4列输出
    3. 使用临时DataFrame避免污染原始数据
    """
    # 输入校验
    required_cols = {'code', 'date', target_col}
    missing_cols = required_cols - set(df.columns)
    if missing_cols:
        raise ValueError(f"缺失必要列: {missing_cols}")
    
    # 创建临时工作DataFrame（仅包含必要列）
    work_df = df[['code', 'date', target_col]].copy()
    
    # 数据预处理
    work_df['date'] = pd.to_datetime(work_df['date'])
    work_df.sort_values(by=['code', 'date'], inplace=True)
    
    # 特征列命名
    suffix = f"_{feature_suffix}" if feature_suffix else ''
    growth_col = f"{target_col}_yoy{suffix}_{window}d"
    pre_growth_col = f"pre_{growth_col}"
    
    # 零值特殊处理 - 添加微小常数避免除零错误
    epsilon = 1e-10
    work_df['_adjusted_target'] = work_df[target_col].replace(0, epsilon) + epsilon
    
    # 滚动基准计算（优化零值处理）
    work_df['_base'] = work_df.groupby('code', group_keys=False)['_adjusted_target'].transform(
        lambda x: x.shift(1).rolling(window, min_periods=1).mean()
    )
    # display(work_df[-30:])
    # 增长率计算（优化零值处理）
    prev_col = work_df.groupby('code')[target_col].shift(1)
    
    if growth_type == 'relative':
        # 核心优化：零值特殊处理
        with np.errstate(divide='ignore', invalid='ignore'):
            work_df[growth_col] = np.where(
                work_df['_base'] <= min_valid_value,
                # 基准接近零时的替代计算
                (work_df[target_col] - prev_col) / (prev_col + epsilon),
                # 正常相对增长率计算
                (work_df[target_col] - work_df['_base']) / work_df['_base']
            )
        # display(work_df[-30:])
    elif growth_type == 'log':
        # 对数增长率同样需要零值处理
        work_df[growth_col] = np.where(
            (work_df['_base'] > min_valid_value) & (work_df[target_col] > min_valid_value),
            np.log(work_df[target_col] / work_df['_base']),
            np.where(
                # 零到非零的跃迁
                (work_df[target_col] > min_valid_value) & (prev_col <= min_valid_value),
                1.0,  # 固定值表示显著增长
                np.nan
            )
        )
    else:
        raise ValueError("growth_type 必须为 'relative' 或 'log'")
    
    # 计算昨日特征
    work_df[pre_growth_col] = work_df.groupby('code')[growth_col].shift(1)
    
    # 仅保留最终需要的4列
    result = work_df[['code', 'date', growth_col, pre_growth_col]].copy()
    
    # 停牌数据处理（如果提供停牌列）
    if trading_col in df.columns:
        # 合并停牌信息
        trading_info = df[['code', 'date', trading_col]].drop_duplicates()
        result = result.merge(trading_info, on=['code', 'date'], how='left')
        result.loc[~result[trading_col], [growth_col, pre_growth_col]] = np.nan
        result.drop(columns=[trading_col], inplace=True)
    
    return result.round(4)  # 提高精度保留小数

    